import cv2
import numpy as np
from pyapriltags import Detector
import json
import time
import threading

class CameraTracker:
    def __init__(self, rtsp_url):
        self.rtsp_url = rtsp_url
        self.detector = Detector(families="tag36h11")
        self.latest_frame = None
        self.latest_map = None
        self.positions = []
        self.running = False
        self.thread = None
        
        # 加载相机参数
        try:
            params = np.load('camera_params.npz')
            self.camera_matrix = params['camera_matrix']
            self.dist_coeffs = params['dist_coeffs']
        except:
            print("警告: 未找到相机参数文件")
            self.camera_matrix = None
            self.dist_coeffs = None
    
    def start_tracking(self):
        """启动跟踪线程"""
        if self.running:
            return {"status": "已在运行中"}
        
        if self.camera_matrix is None:
            return {"error": "请先进行相机标定"}
        
        self.running = True
        self.thread = threading.Thread(target=self._tracking_loop)
        self.thread.daemon = True
        self.thread.start()
        return {"status": "跟踪已启动"}
    
    def stop_tracking(self):
        """停止跟踪"""
        self.running = False
        if self.thread:
            self.thread.join(timeout=2.0)
        return {"status": "跟踪已停止"}
    
    def _tracking_loop(self):
        """跟踪主循环"""
        cap = cv2.VideoCapture(self.rtsp_url)
        if not cap.isOpened():
            print("无法连接摄像头")
            return
        
        # 人员标记配置
        person_tags = {0: "wyx"}
        tag_size = 0.07  # 标记尺寸
        
        # 定义标记的3D点
        object_points = np.array([
            [-tag_size/2, -tag_size/2, 0],
            [tag_size/2, -tag_size/2, 0],
            [tag_size/2, tag_size/2, 0],
            [-tag_size/2, tag_size/2, 0]
        ], dtype=np.float32)
        
        while self.running:
            ret, frame = cap.read()
            if not ret: 
                time.sleep(0.1)
                continue
            
            # 处理帧
            processed_frame, positions = self._process_frame(frame, person_tags, object_points)
            
            # 创建平面图
            map_img = self._create_map(positions)
            
            # 更新共享数据
            self.latest_frame = processed_frame
            self.latest_map = map_img
            self.positions = positions
            
            # 控制处理频率
            time.sleep(0.05)
        
        cap.release()
    
    def _process_frame(self, frame, person_tags, object_points):
        """处理单帧图像"""
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        results = self.detector.detect(gray)
        positions = []
        
        for r in results:
            if r.tag_id in person_tags:
                corners = r.corners.astype(np.float32)
                
                # 计算位姿
                success, rvec, tvec = cv2.solvePnP(
                    object_points, corners, 
                    self.camera_matrix, self.dist_coeffs
                )
                
                if success:
                    position = tvec.flatten()
                    positions.append({
                        "id": r.tag_id,
                        "name": person_tags[r.tag_id],
                        "x": float(position[0]),
                        "y": float(position[1]),
                        "z": float(position[2])
                    })
                    
                    # 绘制边界框
                    corners_int = corners.astype(int)
                    cv2.polylines(frame, [corners_int], True, (0, 255, 0), 2)
                    
                    # 显示位置信息
                    info = f"{person_tags[r.tag_id]}: X={position[0]:.2f}m, Z={position[2]:.2f}m"
                    cv2.putText(frame, info, 
                                (int(corners_int[0][0]), int(corners_int[0][1]) - 20), 
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
        
        return frame, positions
    
    def _create_map(self, positions):
        """创建平面图"""
        map_img = np.zeros((400, 600, 3), dtype=np.uint8) + 50
        cv2.rectangle(map_img, (50, 50), (550, 350), (200, 200, 200), 2)
        
        # 添加刻度
        for i in range(0, 4):
            y_pos = 50 + i * 100
            cv2.line(map_img, (45, y_pos), (55, y_pos), (200, 200, 200), 1)
            cv2.putText(map_img, f"{i}m", (20, y_pos+5), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, (200, 200, 200), 1)
        
        # 绘制人员位置
        for person in positions:
            map_x = int(100 + person['x'] * 100)
            map_y = int(300 - person['z'] * 100)  # Z轴反向
            
            if 50 <= map_x <= 550 and 50 <= map_y <= 350:
                cv2.circle(map_img, (map_x, map_y), 8, (0, 255, 0), -1)
                cv2.putText(map_img, person['name'], 
                            (map_x+10, map_y-10), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
        
        return map_img
    
    def get_latest_data(self):
        """获取最新数据"""
        return {
            "frame": self._encode_image(self.latest_frame) if self.latest_frame is not None else None,
            "map": self._encode_image(self.latest_map) if self.latest_map is not None else None,
            "positions": self.positions
        }
    
    def _encode_image(self, image):
        """将OpenCV图像编码为base64"""
        if image is None:
            return None
        _, buffer = cv2.imencode('.jpg', image)
        return buffer.tobytes()